import math
import os
import sys
from pathlib import Path
import cv2
import numpy as np
import torch

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # yolov5 strongsort root directory
WEIGHTS = ROOT / 'weights'

if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
if str(ROOT / 'yolov5') not in sys.path:
    sys.path.append(str(ROOT / 'yolov5'))  # add yolov5 ROOT to PATH

from yolov5.models.common import DetectMultiBackend
from yolov5.utils.general import check_img_size
from yolov5.utils.torch_utils import select_device
from yolov5.utils.augmentations import letterbox
#prepare for yolo

device = ''
device = select_device(device)
imgsz = (640,640)
yolo_weights,dnn,half = 'weights/best.pt',False,False
model = DetectMultiBackend(yolo_weights, device=device, dnn=dnn, data=None, fp16=half)
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride)  # check image size
augment = False
visualize = False

#nms
conf_thres=0.25
iou_thres=0.45
max_det=1000
agnostic_nms=False
classes=None

#prepare for video
# cudnn.benchmark = True  # set True to speed up constant image size inference

cap1 = cv2.VideoCapture(0)

def next_img(cap1):
    _,img0 = cap1.read()
    # img0 = np.rot90(img0)
    img = letterbox(img0, imgsz, stride=stride)[0]
    img = np.array([img])
    img = img[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW
    img = np.ascontiguousarray(img)

    return img,np.ascontiguousarray(img0)

# print(model)





